package com.ruge.util;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;

import java.util.HashMap;
import java.util.Map;

public class TreeUtil {

    public static JSONArray buildTree(JSONArray deptArray, String idKey, String pidKey) {
        // 使用Map存储所有节点，方便快速查找
        Map<String, JSONObject> nodeMap = new HashMap<>();
        // 根节点列表
        JSONArray rootNodes = new JSONArray();

        // 第一遍遍历：将所有节点存入Map，并初始化children数组
        for (Object obj : deptArray) {
            JSONObject dept = (JSONObject) obj;
            dept.put("children", new JSONArray());
            nodeMap.put(dept.getString(idKey), dept);
        }

        // 第二遍遍历：构建树结构
        for (Object obj : deptArray) {
            JSONObject dept = (JSONObject) obj;
            String code = dept.getString(idKey);
            String flcode = dept.getString(pidKey);

            // 如果父节点存在于Map中，则将当前节点添加为其子节点
            JSONObject parent = nodeMap.get(flcode);
            if (parent != null) {
                JSONArray children = parent.getJSONArray("children");
                children.add(dept);
            } else {
                // 如果找不到父节点，说明是根节点
                rootNodes.add(dept);
            }
        }

        return rootNodes;
    }
}
